import numpy as np
import pandas as pd
from scipy.interpolate import interp1d


# %%
def t_r_unique(t, r):
    df = pd.DataFrame(
        np.transpose(np.concatenate(([t], r), axis=0)),
        columns=["t", "rx", "ry", "rz"],
    )
    # df = df.drop_duplicates()  # full dublicates
    df = df.drop_duplicates(subset=["t"])  # with same value in "t" column
    t = df["t"].values.tolist()
    r = df.iloc[:, 1:].T.values.tolist()
    return t, r


# %%
def interpolate_positions(sat1_sol, sat2_sol, free_flight_single_sol):
    t1 = sat1_sol[0][1:]  # ignore first (zero index) point
    t2 = sat2_sol[0][1:]  # ignore first (zero index) point
    r1 = sat1_sol[1][2:5, 1:]  # ignore first (zero index) point
    r2 = sat2_sol[1][2:5, 1:]  # ignore first (zero index) point
    t_free = free_flight_single_sol.t
    r_free = free_flight_single_sol.y[0:3]
    if t1[-1] < t2[-1]:
        t1 = np.concatenate((t1, t_free), axis=0)
        r1 = np.concatenate((r1, r_free), axis=1)
    else:
        t2 = np.concatenate((t2, t_free), axis=0)
        r2 = np.concatenate((r2, r_free), axis=1)

    t1, r1 = t_r_unique(t1, r1)
    t2, r2 = t_r_unique(t2, r2)
    r1_interp = interp1d(t1, r1, kind="cubic")
    r2_interp = interp1d(t2, r2, kind="cubic")
    return r1_interp, r2_interp
